#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File    : mngs_split_sam.py
@Author  : Bing Liang
@Email   : believer19940901@gmail.com
@Date    : 2025/9/25 15:21
@Description : 从总sam文件中拆分不同样本的比对信息
@Usage   : 标准输入流 | python3 mngs_split_sam.py 输出目录 物种种类
"""
import sys
from pathlib import Path


def main():

    out_dir = Path(sys.argv[1]).absolute()
    cate = sys.argv[2].strip()

    file_handles = {}

    try:
        for line in sys.stdin:
            line = line.strip()
            if not line or line.startswith("@"):  # 跳过空行和注释
                continue

            # 跳过未必对上的序列
            if line.split("\t")[1] not in  ["0", "16"]:
                continue

            # 分离样本名和比对信息（假设只按第一个下划线分割）
            if "_" not in line:
                continue
            sample, info = line.split("_", 1)

            # 打开对应样本的文件句柄
            if sample not in file_handles:
                sample_dir = out_dir / sample
                sample_dir.mkdir(parents=True, exist_ok=True)
                fh = open(sample_dir / f"{cate}.sam_tmp", "w", encoding="utf-8")
                file_handles[sample] = fh

            # 写入比对信息
            file_handles[sample].write(info + "\n")

    finally:
        # 确保文件关闭
        for fh in file_handles.values():
            fh.close()


if __name__ == "__main__":
    main()
